// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    file_format::{
        AddressPoolIndex, ByteArrayPoolIndex, Bytecode, CodeOffset, CodeUnit, FieldDefinitionIndex,
        FunctionDefinition, FunctionHandle, FunctionHandleIndex, FunctionSignature,
        FunctionSignatureIndex, LocalIndex, LocalsSignature, LocalsSignatureIndex,
        ModuleHandleIndex, StringPoolIndex, StructDefinitionIndex, TableIndex,
    },
    proptest_types::signature::{FunctionSignatureGen, SignatureTokenGen},
};
use proptest::{
    collection::{vec, SizeRange},
    prelude::*,
    sample::{select, Index as PropIndex},
};

/// Represents state required to materialize final data structures for function definitions.
#[derive(Debug)]
pub struct FnDefnMaterializeState {
    pub struct_handles_len: usize,
    pub address_pool_len: usize,
    pub string_pool_len: usize,
    pub byte_array_pool_len: usize,
    pub type_signatures_len: usize,
    pub field_defs_len: usize,
    pub struct_defs_len: usize,
    pub function_defs_len: usize,
    // This is the final length of function_handles, after all the definitions add their own
    // handles.
    pub function_handles_len: usize,
    // These get mutated by `FunctionDefinitionGen`.
    pub function_signatures: Vec<FunctionSignature>,
    pub locals_signatures: Vec<LocalsSignature>,
    pub function_handles: Vec<FunctionHandle>,
}

impl FnDefnMaterializeState {
    #[inline]
    fn add_function_signature(&mut self, sig: FunctionSignature) -> FunctionSignatureIndex {
        self.function_signatures.push(sig);
        FunctionSignatureIndex::new((self.function_signatures.len() - 1) as TableIndex)
    }

    #[inline]
    fn add_locals_signature(&mut self, sig: LocalsSignature) -> LocalsSignatureIndex {
        self.locals_signatures.push(sig);
        LocalsSignatureIndex::new((self.locals_signatures.len() - 1) as TableIndex)
    }

    #[inline]
    fn add_function_handle(&mut self, handle: FunctionHandle) -> FunctionHandleIndex {
        self.function_handles.push(handle);
        FunctionHandleIndex::new((self.function_handles.len() - 1) as TableIndex)
    }
}

#[derive(Clone, Debug)]
pub struct FunctionDefinitionGen {
    name: PropIndex,
    signature: FunctionSignatureGen,
    is_public: bool,
    code: CodeUnitGen,
}

impl FunctionDefinitionGen {
    pub fn strategy(
        return_count: impl Into<SizeRange>,
        arg_count: impl Into<SizeRange>,
        code_len: impl Into<SizeRange>,
    ) -> impl Strategy<Value = Self> {
        let return_count = return_count.into();
        let arg_count = arg_count.into();
        (
            any::<PropIndex>(),
            FunctionSignatureGen::strategy(return_count.clone(), arg_count.clone()),
            any::<bool>(),
            CodeUnitGen::strategy(arg_count, code_len),
        )
            .prop_map(|(name, signature, is_public, code)| Self {
                name,
                signature,
                is_public,
                code,
            })
    }

    pub fn materialize(self, state: &mut FnDefnMaterializeState) -> FunctionDefinition {
        let signature = self.signature.materialize(state.struct_handles_len);

        let handle = FunctionHandle {
            // 0 represents the current module
            module: ModuleHandleIndex::new(0),
            // XXX need to guarantee uniqueness of names?
            name: StringPoolIndex::new(self.name.index(state.string_pool_len) as TableIndex),
            signature: state.add_function_signature(signature),
        };
        let function_handle = state.add_function_handle(handle);

        FunctionDefinition {
            function: function_handle,
            // XXX is this even correct?
            flags: if self.is_public {
                CodeUnit::PUBLIC
            } else {
                // No qualifiers.
                0
            },
            code: self.code.materialize(state),
        }
    }
}

#[derive(Clone, Debug)]
struct CodeUnitGen {
    locals_signature: Vec<SignatureTokenGen>,
    code: Vec<BytecodeGen>,
}

impl CodeUnitGen {
    fn strategy(
        arg_count: impl Into<SizeRange>,
        code_len: impl Into<SizeRange>,
    ) -> impl Strategy<Value = Self> {
        (
            vec(SignatureTokenGen::strategy(), arg_count),
            vec(BytecodeGen::garbage_strategy(), code_len),
        )
            .prop_map(|(locals_signature, code)| Self {
                locals_signature,
                code,
            })
    }

    fn materialize(self, state: &mut FnDefnMaterializeState) -> CodeUnit {
        let locals_signature = LocalsSignature(
            self.locals_signature
                .into_iter()
                .map(|sig| sig.materialize(state.struct_handles_len))
                .collect(),
        );

        // Not all bytecodes will be successfully materialized -- count how many will.
        let code_len = self
            .code
            .iter()
            .filter(|code| code.will_materialize(state, &locals_signature))
            .count();

        let code = self
            .code
            .into_iter()
            .filter_map(|code| code.materialize(state, code_len, &locals_signature))
            .collect();

        CodeUnit {
            max_stack_size: 0,
            locals: state.add_locals_signature(locals_signature),
            // XXX actually generate code
            code,
        }
    }
}

#[derive(Clone, Debug)]
enum BytecodeGen {
    // "Simple" means this doesn't refer to any other indexes.
    Simple(Bytecode),
    // All of these refer to other indexes.
    LdAddr(PropIndex),
    LdStr(PropIndex),
    LdByteArray(PropIndex),
    BorrowField(PropIndex),
    Call(PropIndex),
    Pack(PropIndex),
    Unpack(PropIndex),
    Exists(PropIndex),
    BorrowGlobal(PropIndex),
    MoveFrom(PropIndex),
    MoveToSender(PropIndex),
    BrTrue(PropIndex),
    BrFalse(PropIndex),
    Branch(PropIndex),
    CopyLoc(PropIndex),
    MoveLoc(PropIndex),
    StLoc(PropIndex),
    BorrowLoc(PropIndex),
}

impl BytecodeGen {
    // This just generates nonsensical bytecodes. This will be cleaned up later as the generation
    // model is refined.
    fn garbage_strategy() -> impl Strategy<Value = Self> {
        use BytecodeGen::*;

        prop_oneof![
            Self::simple_bytecode_strategy().prop_map(Simple),
            any::<PropIndex>().prop_map(LdAddr),
            any::<PropIndex>().prop_map(LdStr),
            any::<PropIndex>().prop_map(LdByteArray),
            any::<PropIndex>().prop_map(BorrowField),
            any::<PropIndex>().prop_map(Call),
            any::<PropIndex>().prop_map(Pack),
            any::<PropIndex>().prop_map(Unpack),
            any::<PropIndex>().prop_map(Exists),
            any::<PropIndex>().prop_map(BorrowGlobal),
            any::<PropIndex>().prop_map(MoveFrom),
            any::<PropIndex>().prop_map(MoveToSender),
            any::<PropIndex>().prop_map(BrTrue),
            any::<PropIndex>().prop_map(BrFalse),
            any::<PropIndex>().prop_map(Branch),
            any::<PropIndex>().prop_map(CopyLoc),
            any::<PropIndex>().prop_map(MoveLoc),
            any::<PropIndex>().prop_map(StLoc),
            any::<PropIndex>().prop_map(BorrowLoc),
        ]
    }

    /// Whether this code will be materialized into a Some(bytecode).
    fn will_materialize(
        &self,
        state: &FnDefnMaterializeState,
        locals_signature: &LocalsSignature,
    ) -> bool {
        // This method should remain in sync with the `None` below.
        use BytecodeGen::*;

        match self {
            BorrowField(_) => state.field_defs_len != 0,
            CopyLoc(_) | MoveLoc(_) | StLoc(_) | BorrowLoc(_) => !locals_signature.is_empty(),
            _ => true,
        }
    }

    fn materialize(
        self,
        state: &FnDefnMaterializeState,
        code_len: usize,
        locals_signature: &LocalsSignature,
    ) -> Option<Bytecode> {
        // This method returns an Option<Bytecode> because some bytecodes cannot be represented if
        // some tables are empty.
        //
        // Once more sensible function bodies are generated this will probably have to start using
        // prop_flat_map anyway, so revisit this then.

        let bytecode = match self {
            BytecodeGen::Simple(bytecode) => bytecode,
            BytecodeGen::LdAddr(idx) => Bytecode::LdAddr(AddressPoolIndex::new(
                idx.index(state.address_pool_len) as TableIndex,
            )),
            BytecodeGen::LdStr(idx) => Bytecode::LdStr(StringPoolIndex::new(
                idx.index(state.string_pool_len) as TableIndex,
            )),
            BytecodeGen::LdByteArray(idx) => Bytecode::LdByteArray(ByteArrayPoolIndex::new(
                idx.index(state.byte_array_pool_len) as TableIndex,
            )),
            BytecodeGen::BorrowField(idx) => {
                // Again, once meaningful bytecodes are generated this won't actually be a
                // possibility since it would be impossible to load a field from a struct that
                // doesn't have any.
                if state.field_defs_len == 0 {
                    return None;
                }
                Bytecode::BorrowField(FieldDefinitionIndex::new(
                    idx.index(state.field_defs_len) as TableIndex
                ))
            }
            BytecodeGen::Call(idx) => Bytecode::Call(FunctionHandleIndex::new(
                idx.index(state.function_handles_len) as TableIndex,
            )),
            BytecodeGen::Pack(idx) => Bytecode::Pack(StructDefinitionIndex::new(
                idx.index(state.struct_defs_len) as TableIndex,
            )),
            BytecodeGen::Unpack(idx) => Bytecode::Unpack(StructDefinitionIndex::new(
                idx.index(state.struct_defs_len) as TableIndex,
            )),
            BytecodeGen::Exists(idx) => Bytecode::Exists(StructDefinitionIndex::new(
                idx.index(state.struct_defs_len) as TableIndex,
            )),
            BytecodeGen::BorrowGlobal(idx) => Bytecode::BorrowGlobal(StructDefinitionIndex::new(
                idx.index(state.struct_defs_len) as TableIndex,
            )),
            BytecodeGen::MoveFrom(idx) => Bytecode::MoveFrom(StructDefinitionIndex::new(
                idx.index(state.struct_defs_len) as TableIndex,
            )),
            BytecodeGen::MoveToSender(idx) => Bytecode::MoveToSender(StructDefinitionIndex::new(
                idx.index(state.struct_defs_len) as TableIndex,
            )),
            BytecodeGen::BrTrue(idx) => Bytecode::BrTrue(idx.index(code_len) as CodeOffset),
            BytecodeGen::BrFalse(idx) => Bytecode::BrFalse(idx.index(code_len) as CodeOffset),
            BytecodeGen::Branch(idx) => Bytecode::Branch(idx.index(code_len) as CodeOffset),
            BytecodeGen::CopyLoc(idx) => {
                if locals_signature.is_empty() {
                    return None;
                }
                Bytecode::CopyLoc(idx.index(locals_signature.len()) as LocalIndex)
            }
            BytecodeGen::MoveLoc(idx) => {
                if locals_signature.is_empty() {
                    return None;
                }
                Bytecode::MoveLoc(idx.index(locals_signature.len()) as LocalIndex)
            }
            BytecodeGen::StLoc(idx) => {
                if locals_signature.is_empty() {
                    return None;
                }
                Bytecode::StLoc(idx.index(locals_signature.len()) as LocalIndex)
            }
            BytecodeGen::BorrowLoc(idx) => {
                if locals_signature.is_empty() {
                    return None;
                }
                Bytecode::BorrowLoc(idx.index(locals_signature.len()) as LocalIndex)
            }
        };

        Some(bytecode)
    }

    fn simple_bytecode_strategy() -> impl Strategy<Value = Bytecode> {
        prop_oneof![
            // The numbers are relative weights, somewhat arbitrarily picked.
            9 => Self::just_bytecode_strategy(),
            1 => any::<u64>().prop_map(Bytecode::LdConst),
        ]
    }

    fn just_bytecode_strategy() -> impl Strategy<Value = Bytecode> {
        use Bytecode::*;

        static JUST_BYTECODES: &[Bytecode] = &[
            FreezeRef,
            ReleaseRef,
            Pop,
            Ret,
            LdTrue,
            LdFalse,
            ReadRef,
            WriteRef,
            Add,
            Sub,
            Mul,
            Mod,
            Div,
            BitOr,
            BitAnd,
            Xor,
            Or,
            And,
            Eq,
            Neq,
            Lt,
            Gt,
            Le,
            Ge,
            Assert,
            GetTxnGasUnitPrice,
            GetTxnMaxGasUnits,
            GetTxnSenderAddress,
            CreateAccount,
            EmitEvent,
            GetTxnSequenceNumber,
            GetTxnPublicKey,
        ];
        select(JUST_BYTECODES)
    }
}
